"""
author: liuxu
date: 2024.12.3
description: 执行ssh命令模块
"""
import os
import sys
import re
import time
import subprocess
import select
import tempfile
import threading
import traceback
import pdb

from common.logger_ctrl import mylogger
from typing import Generator, Tuple, Dict, Any, Optional, List, Union, Callable


glo_command = 0 #用于全局执行ssh命令计数

# 这些会打印在stdout，需特殊处理
excluded_err_output = {"HIPPROFILE-trace_hsa_call: hsa_support OnLoad",
                       "Could not open /var/log/hylog/.",
                       "HIPPROFILE-trace_hsa_call: unknown AmdExtTable version 1.624",
                       "HIPPROFILE-trace_hsa_call: unknown AmdExtTable version 1.648",
                       "Using the old table 6.3.10, previously size 608 now 324",
                        "Using the old table 6.3.10, previously size 608 now 504",
                        "Using the old table 6.3.10, previously size 608 now 624",
                        "Using the old table 6.3.10, previously size 608 now 648"
                       }

def exec_cmd(cmd, *args, **kwargs) -> Tuple[bool, str]:
    """
    执行带参数的linux命令
    :param cmd:
    :param args:
    :param kwargs:
    :return: (命令执行是否成功， 执行输出结果)
    """
    try:
        global glo_command
        glo_command += 1
        # cwd = kwargs.get("cwd", TMP_DIR)
        # 分割需要执行的命令行参数
        mylogger.info(f"{'>' * 20}Executing command {glo_command}: {cmd} {' '.join(args)} {'<' * 20}")
        new_cmd = cmd.split(" ") if isinstance(cmd, str) else cmd
        # 设置subprocess.run默认参数
        kwargs.setdefault('timeout', 30)  # 设置命令默认超时时间为30s

        # 设置check=True，返回非0错误码则抛出异常
        ret = subprocess.run([*new_cmd, *args], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, **kwargs)
        output = ret.stdout.strip()
        error = ret.stderr.strip()
        # 打印标准输出和错误输出
        mylogger.info(f"{'>' * 20}command {glo_command} output:\n{output}")
        output = re.sub(r'\x1b\[.*?m', '', output)
        error = re.sub(r'\x1b\[.*?m', '', error)
        if error:
            mylogger.error(f"{'>' * 20}command {glo_command} stderr:\n{error}")
            # 过滤预期内的stderr
            if set(error.split("\n")).issubset(excluded_err_output):
                return True, output
            return False, output + error
        return True, output
    except subprocess.CalledProcessError as e:
        mylogger.error(f"Command execution failed, error message: {e}")
        e_output = e.stdout.strip()
        e_err = e.stderr.strip()
        mylogger.debug(f"Exception stdout:\n{e_output}")
        mylogger.debug(f"Exception stderr:\n{e_err}")
        return False, e_output + e_err  or ""
    except subprocess.TimeoutExpired as e:
        mylogger.error(f"Command execution timeout, error message: {e}")
        return False, e.stdout.strip() or ""
    except Exception as e:
        mylogger.error(f"Command execution failed, error message: {e}\n{traceback.format_exc()}")
    return False,""

def interactive_exec_cmd(**kwargs):
    """
    执行ssh交互式命令
    :return:
    """
    try:
        # 启动bash子进程
        prompt = "AUTO_CMD_END_PROMPT"
        process = subprocess.Popen(
            ["/bin/bash", "--norc", "--noprofile", "-i"],  # 交互模式
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=0,
            universal_newlines=True,
            env={"PS1": prompt},  # 设置命令提示符号
            **kwargs
        )
        mylogger.debug(f"Current subprocess popen PID: {process.pid}")
    except Exception as e:
        mylogger.error(f"SSHException!\n{e}\n{traceback.format_exc()}")
        yield False, "Failed to start subprocess"
        return

    stdout_ = ""
    stderr_ = ""

    try:
        # 初始启动生成器
        command, timeout = yield stdout_, stderr_

        while command is not None:
            global glo_command
            glo_command += 1
            mylogger.info(f"the {glo_command} ssh command stdin: {command}")

            # 发送命令
            process.stdin.write(command + "\n")
            process.stdin.flush()

            # 准备读取输出
            stdout_buffer = ""
            stderr_buffer = ""
            start_time = time.time()

            while True:
                # 计算剩余时间
                if time.time() - start_time > timeout:
                    # mylogger.warning("Command timed out")
                    break

                # 使用select监控输出管道
                rlist, _, _ = select.select(
                    [process.stdout, process.stderr],
                    [],
                    [],
                    0.1 # 等待最多0.1秒或剩余时间
                )

                # 处理标准输出
                if process.stdout in rlist:
                    chunk = os.read(process.stdout.fileno(), 4096)
                    if chunk:
                        chunk = chunk.decode('utf-8', errors='replace')
                        stdout_buffer += chunk
                        # # 实时输出（可选）
                        # sys.stdout.write(chunk)
                        # sys.stdout.flush()

                # 处理错误输出
                if process.stderr in rlist:
                    chunk = os.read(process.stderr.fileno(), 4096)
                    if chunk:
                        chunk = chunk.decode('utf-8', errors='replace')
                        stderr_buffer += chunk
                        # # 实时输出（可选）
                        # sys.stderr.write(chunk)
                        # sys.stderr.flush()

                # 检查终止条件：找到提示符或超时
                if prompt in stdout_buffer:
                    # 提取提示符前的输出
                    prompt_index = stdout_buffer.rfind(prompt)
                    if prompt_index != -1:
                        stdout_buffer = stdout_buffer[:prompt_index].strip()
                    break

            # 更新最终输出
            stdout_ = stdout_buffer.strip()
            stderr_ = stderr_buffer.strip()

            mylogger.info(f"stdout:\n{stdout_}")
            mylogger.warning(f"stderr:\n{stderr_}")

            # 等待下一个命令
            command, timeout = yield stdout_, stderr_

    except Exception as e:
        # mylogger.error(f"Exception occurred: {e}\n{traceback.format_exc()}")
        yield f"Error: {str(e)}", traceback.format_exc()
    finally:
        # 清理资源
        process.stdin.close()
        process.terminate()
        process.wait()
        mylogger.info(f"Subprocess {process.pid} terminated")
        yield "Process exited 0", ""


class InteractiveExecutor:
    """
    通用命令行交互式执行器
    支持任何命令行工具的自动化交互测试
    """

    def __init__(
        self,
        command: Union[str, List[str]],
        prompt: Union[str, re.Pattern, Callable[[str], bool]],
        bufsize: int = 0,
        timeout: int = 10,
        env: Optional[Dict[str, str]] = None,
        encoding: str = "utf-8",
        auto_start: bool = True
    ):
        """
        初始化交互式执行器

        :param command: 要执行的命令（字符串或列表）
        :param prompt: 初始提示符（字符串、正则或检测函数）
        :param end_prompt: 结束提示符（字符串、正则或检测函数）
        :param bufsize: 设置缓冲区大小
                0：无缓冲（数据立即传输）
                1：行缓冲（遇到换行符时刷新）
                -1：系统默认（通常为全缓冲）
        :param timeout: 默认命令超时时间（秒）
        :param env: 环境变量字典
        :param encoding: 输出编码
        :param auto_start: 是否自动启动交互进程
        """
        self.command = command.split() if isinstance(command, str) else command
        self.prompt = prompt
        self.default_timeout = timeout
        self.bufsize = bufsize
        self.env = env or {}
        self.encoding = encoding
        self.process = None
        self.is_running = False
        self.lock = threading.Lock()

        # 编译提示符检测函数
        self.check_prompt_func = self._compile_prompt_func(prompt)

        # 设置环境变量
        self.full_env = os.environ.copy()
        self.full_env.update(self.env)

        if auto_start:
            self.start()
        # mylogger.debug(f"cur env:\n{self.full_env}")

    @staticmethod
    def get_glo_cmd_cnt(increase=True):
        """使用全局变量（命令计数器）"""
        global glo_command
        if increase:
            glo_command += 1
        return glo_command


    @staticmethod
    def _compile_prompt_func(prompt) -> Callable[[str], bool]:
        """编译提示符检测函数"""
        if isinstance(prompt, str):
            pattern = re.compile(re.escape(prompt))
            return lambda s: pattern.search(s) is not None
        elif isinstance(prompt, re.Pattern):
            return lambda s: prompt.search(s) is not None
        elif callable(prompt):
            return prompt
        else:
            raise ValueError("Invalid prompt type")

    def start(self) -> bool:
        """启动交互式测试进程"""
        if self.is_running:
            return True
        try:
            self.process = subprocess.Popen(
                self.command,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                # text=True,  # text=True自动将输入（stdin）转化为byte，输出（stdout/stderr)转化为str；Python3.6及以下使用universal_newlines=True替代
                text=False,  # 使用二进制模式，手动处理编码
                bufsize=self.bufsize,
                env=self.full_env
            )
            self.is_running = True

            # 等待初始提示
            mylogger.info(f"the {self.get_glo_cmd_cnt()} ssh command stdin: {self.command}")
            out, err = self._read_until_prompt(self.check_prompt_func, self.default_timeout)
            mylogger.info(f"stdout:\n{out}")
            if err:
                mylogger.error(f"stderr:\n{err}")
            mylogger.debug(f"Succeed to start subprocess, current popen PID: {self.process.pid}")
            return True

        except Exception as e:
            raise RuntimeError(f"Process start failed: {e}\n{traceback.format_exc()}")

    def _read_until_prompt(
            self,
            prompt_func,
            timeout: int
    ) -> Tuple[str, str]:
        """
        读取输出直到满足条件
        :param prompt_func: Callable[[str], bool],
        :param timeout: 超时时间（秒）
        :return: (stdout, stderr)
        """
        stdout_str = ""
        stderr_str = ""
        start_time = time.time()

        while True:
            # 检查命令是否超时
            if time.time() - start_time > timeout:
                mylogger.error(f"Timeout waiting for prompt '{prompt_func.__closure__[0].cell_contents.pattern}' after {timeout}s")
                break

            # 检查进程是否已结束
            if self.process.poll() is not None:  # 交互式模式输入退出字符，会关闭进程
                mylogger.debug("Process has ended during reading")
                break
            # pdb.set_trace()  # debug

            # 使用select.select()和os.read()监控管道,非阻塞读取子进程输出
            rlist, _, _ = select.select(
                [self.process.stdout, self.process.stderr],  # 等待读取的文件描述符
                [], # 等待写入的文件描述符
                [],  # 等待写入的文件描述符
                0.1  # 超时时间（秒）
            )

            # 读取标准输出
            if self.process.stdout in rlist:
                chunk = os.read(self.process.stdout.fileno(), 4096)
                if chunk:
                    # 解码输出(errors='replace'：遇到无法解码的字符时用?替代，避免抛出异常)
                    stdout_str += chunk.decode(encoding=self.encoding, errors='replace')

            # 读取错误输出
            if self.process.stderr in rlist:
                chunk = os.read(self.process.stderr.fileno(), 4096)
                if chunk:
                    stderr_str += chunk.decode(encoding=self.encoding, errors='replace')

            # 检查提示条件
            if prompt_func(stdout_str) or prompt_func(stderr_str):
                break

        return stdout_str, stderr_str

    def execute(
            self,
            command: str,
            prompt: Optional[Union[str, re.Pattern, Callable[[str], bool]]] = None,
            timeout: Optional[int] = None
    ) -> Tuple[str, str]:
        """
        执行单个命令
        :param command: 要执行的命令
        :param prompt： 命令执行结束符
        :param timeout: 超时时间（秒）
        :return: (stdout, stderr)
        """
        with self.lock:
            if not self.is_running or not self.process:
                raise RuntimeError("Process is not running")

            timeout = timeout or self.default_timeout

            try:
                # 发送命令
                self.process.stdin.write((command + "\n").encode(self.encoding))
                self.process.stdin.flush()
                mylogger.info(f"the {self.get_glo_cmd_cnt()} ssh command stdin:\n {command}")

                # 等待提示符出现（命令完成）
                if prompt is None:
                    prompt_func = self.check_prompt_func
                else:
                    prompt_func = self._compile_prompt_func(prompt)
                stdout_str, stderr_str = self._read_until_prompt(prompt_func, timeout)
                # pdb.set_trace()
                mylogger.info(f"the {self.get_glo_cmd_cnt(False)} ssh command stdout:\n{stdout_str}")
                if stderr_str:
                    mylogger.warning(f"the {self.get_glo_cmd_cnt(False)} ssh command stderr:\n{stderr_str}")
                return stdout_str, stderr_str

            except Exception as e:
                return f"Error: {str(e)}", traceback.format_exc()

    def execute_script(
            self,
            script: Union[str, List[str]],
            timeout: Optional[int] = None
    ) -> List[Tuple[str, str]]:
        """
        执行脚本（一系列命令）
        :param script: 脚本内容（字符串或行列表）
        :param timeout: 每个命令的超时时间
        :return: 结果元组列表 (stdout, stderr)
        """
        if isinstance(script, str):
            commands = [cmd.strip() for cmd in script.splitlines() if cmd.strip()]
        else:
            commands = [cmd.strip() for cmd in script if cmd.strip()]

        results = []

        for cmd in commands:
            # 跳过注释
            if cmd.startswith("#"):
                continue

            result = self.execute(cmd, timeout=timeout)
            results.append(result)

        return results

    def stop(self, exit_cmd: str="exit", timeout: int = 1) -> None:
        """
        停止/退出测试进程
        :param exit_cmd:
        :param timeout:
        :return:
        """
        if not self.is_running or not self.process:
            return

        try:
            if "gdb" in self.prompt:
                exit_cmd = "quit"
            elif ">>>" in self.prompt:
                exit_cmd = "exit()"
            # 尝试正常退出
            self.execute(exit_cmd, timeout=timeout)

            # 等待进程结束
            try:
                self.process.wait(timeout=timeout)
                mylogger.debug(f"Succeed to exit subprocess (PID: {self.process.pid})")
            except subprocess.TimeoutExpired:
                self.process.terminate()
                try:
                    self.process.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    self.process.kill()
                    self.process.wait()

        finally:
            # 确保清理资源
            if self.process:
                for stream in [self.process.stdin, self.process.stdout, self.process.stderr]:
                    if stream:
                        try:
                            stream.close()
                        except Exception as e:
                            mylogger.debug(f"Error closing stream: {e}")  # 记录具体错误

            self.process = None
            self.is_running = False

    def __enter__(self):
        """上下文管理器入口"""
        if not self.is_running:
            self.start()
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """上下文管理器出口"""
        self.stop()


if __name__ == '__main__':
    exec_cmd("hy-smi --showproductname | grep 'Card Series' | awk '{print $NF}'")
    exit(0)
    def interactive_exec_cmd_test():
        generator = interactive_exec_cmd()
        next(generator)
        for cmd in [ "hipprof -h", "pwd", "ls", "touch tmp", None]:
            info, err = generator.send((cmd, 1))
            print(info, err)

    def InteractiveExecutor_test():
        with InteractiveExecutor(
            command=["python", "-i", "-q"],  # 启动安静模式的交互式Python
            prompt=">>> ",  # Python提示符
            timeout=5
        ) as py:
            # 执行Python命令
            stdout, _ = py.execute("print('Hello, World!')")
            stdout, _ = py.execute("print(1+1)")

            # 执行脚本
            script = """
            # Python脚本
            import math
            result = math.sqrt(25)
            print(f"Result: {result}")
            """
            results = py.execute_script(script)

            # 检查结果
            last_output = results[-1][0]
            assert "Result: 5.0" in last_output
            py.execute(command="exit()")

        with InteractiveExecutor(
            command=["hipgdb", "./exe/matrix"],  # 启动安静模式的交互式Python
            prompt="(gdb)",  # Python提示符
            timeout=5
        ) as gdb:
            gdb.execute("c")
            gdb.execute("break matrixTranspose")
            gdb.execute("run")
            gdb.execute("c")
            gdb.execute(command="quit")


        # 通过subprocess默认会关闭终端提示符
        # 修改bash启动方式（设置终端类型，确保Bash在非终端环境下仍显示提示符）
        with InteractiveExecutor(
            command=["/bin/bash", "--norc", "-i"],  # --norc跳过配置文件，-i强制交互模式
            prompt=r"(liuxu_auto_test) # ",  # 匹配静态提示符
            timeout=5,
            env={
                "PS1": r"(liuxu_auto_test) # ",  # 静态提示符
                "TERM": "dumb"  # 设置终端类型
                }

        ) as gdb:
            gdb.execute("ls")
            gdb.execute("pwd")
            gdb.execute("exit")

    InteractiveExecutor_test()

